import numpy as np
import paddle
from paddle import nn
from paddle.fluid.layers import detection
import paddle.nn.functional as F
from paddle import optimizer, metric, callbacks
from YOLO.BackBone import DarkNet53_conv_body, ConvBNLayer


class YoloDetectionBlock(paddle.nn.Layer):
    """
    使用多层卷积和BN提取特征
    """
    def __init__(self,ch_in,ch_out,is_test=True):
        super(YoloDetectionBlock, self).__init__()

        assert ch_out % 2 == 0, \
            "channel {} cannot be divided by 2".format(ch_out)

        self.conv0 = ConvBNLayer(
            ch_in=ch_in,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0)
        self.conv1 = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1)
        self.conv2 = ConvBNLayer(
            ch_in=ch_out*2,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0)
        self.conv3 = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1)
        self.route = ConvBNLayer(
            ch_in=ch_out*2,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0)
        self.tip = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1)

    def forward(self, inputs):
        out = self.conv0(inputs)
        out = self.conv1(out)
        out = self.conv2(out)
        out = self.conv3(out)
        route = self.route(out)
        tip = self.tip(route)
        return route, tip

class vehicleDetection(nn.Layer):
    def __init__(self):
        super(vehicleDetection, self).__init__()
        self.backbone = DarkNet53_conv_body()
        self.yoloblock = YoloDetectionBlock(1024, 512)
        # 最后一个预测卷积层，通道为 锚点数量 * 每个锚点预测目标数量（xcenter, ycenter, width_bias, height_bias）
        self.conv2d_pred = nn.Conv2D(1024, 5 * 3, 1)
    
    def forward(self, x):
        x = paddle.transpose(x, [0, 3, 1, 2])
        # x 应该为 NCHW 格式
        C0, C1, C2 = self.backbone(x)
        _, tip = self.yoloblock(C0)
        y = self.conv2d_pred(tip)
        # [batch_size, 5 * 3, num_rows, num_cols]
        return y

class YoloLoss(nn.Layer):
    """
    Yolo 检测专用损失函数。
    标签中的一二数字表示纵、横坐标偏置。
    三四数字表示高和宽的缩放比例。
    最后一个数字表示是否包含物体。
    """

    def __init__(self, name=None):
        super(YoloLoss, self).__init__()
        self.name = name
    
    def forward(self, input, label):
        # input 和 label 的形状：[batch_size, 5 * 3, 20, 10]
        # NOTE num_anchors 在这里固定为3了

        reshaped_input = paddle.reshape(input, (-1, 3, 5, input.shape[2], input.shape[3]))
        reshaped_label = paddle.cast(label, 'float32')
        # [batch_size, num_anchors, 5, height, width]
        pred_objectness = reshaped_input[:, :, -1, :, :]
        label_objectness = reshaped_label[:, :, -1, :, :]
        loss_objectness = F.binary_cross_entropy_with_logits(pred_objectness, label_objectness, reduction='none')
        
        # 标注位置是正样本的
        pos_objectness = label_objectness > 0
        pos_samples = paddle.cast(pos_objectness, 'float32')
        pos_samples.stop_gradient = True

        loss_objectness = loss_objectness * (1 + pos_samples * 100)

        # 预测相关位置信息
        tx = reshaped_input[:, :, 0, :, :]
        ty = reshaped_input[:, :, 1, :, :]
        tw = reshaped_input[:, :, 2, :, :]
        th = reshaped_input[:, :, 3, :, :]

        # 从 reshaped_label 中取出各个位置坐标的标签
        dx_label = reshaped_label[:, :, 0, :, :]
        dy_label = reshaped_label[:, :, 1, :, :]
        tw_label = reshaped_label[:, :, 2, :, :]
        th_label = reshaped_label[:, :, 3, :, :]

        # 构建损失函数
        loss_location_x = F.binary_cross_entropy_with_logits(tx, dx_label, reduction="none")
        loss_location_y = F.binary_cross_entropy_with_logits(ty, dy_label, reduction="none")
        loss_location_w = paddle.abs(tw - tw_label)
        loss_location_h = paddle.abs(th - th_label)

        # 总的位置损失函数
        loss_location = loss_location_x + loss_location_y + loss_location_h + loss_location_w

        # TODO scales

        # 只计算正样本的位置损失函数
        loss_location = loss_location * pos_samples

        total_loss = loss_objectness + loss_location
        # 对所有预测框的loss进行求和
        total_loss = paddle.sum(total_loss, axis=[1,2,3])
        # 对所有样本求平均
        total_loss = paddle.mean(total_loss)

        return total_loss

def build_model(mode='train', path=None):
    detector = vehicleDetection()
    if path:
        detector.set_state_dict(paddle.load(path))
    if mode == 'train':
        detector = paddle.Model(detector)
        detector.prepare(
            optimizer.Adam(0.001, parameters=detector.parameters()),
            YoloLoss()
        )
    return detector

if __name__ == '__main__':
    a = paddle.to_tensor(np.random.rand(5, 3, 640, 320), dtype='float32')
    # backbone = DarkNet53_conv_body()
    # c0, c1, c2 = backbone(a)
    # c0: [5, 1024, 20, 10]

    model = build_model()
    y = model(a)
    print(y.shape)
    # [5, 15, 20, 10]
    print('done')